package com.azure.code.graph.unionFind;

/**
 * 对UF的优化 使 构造树平衡化 优化成logN
 */
public class UFBalance {
    // 连通量
    private int count;
    // 节点的根
    private final int[] parent;

    // 新增一个数组记录数的重量
    private final int[] size;

    public UFBalance(int n) {
        // 一开始互不连通
        this.count = n;
        parent = new int[n];
        size = new int[n];
        // 最初每颗树只有一个节点
        // 重量应该初始化 1
        for (int i = 0; i < n; i++) {
            parent[i] = i;
            size[i] = 1;
        }
    }
    // 查询节点x的根
    public int find(int x){
        while (parent[x]!=x)
            x = parent[x];
        return x;
    }

    /**
     * 将节点p和q连通
     *
     *  size[3] = 5 表示 节点为3的根节点那棵树，总共有5个节点
     *  union 将树平衡化
     */
    public void union(int p,int q){
        int rootP = find(p);
        int rootQ= find(q);
        if(rootP==rootQ)
            return;
//        // 合并两棵树
//        parent[rootP] = rootQ;
        // 小树连接到大树下面，较平衡
        if(size[rootP]>size[rootQ]){
            parent[rootQ] = rootP;
            size[rootP] += size[rootQ];
        }else {
            parent[rootP] = rootQ;
            size[rootQ] += size[rootP];
        }
        // 连通量减一
        count--;
    }

    // 判断节点p和节点q是否连通
    public boolean connected(int p, int q){
        int rootP = find(p);
        int rootQ = find(q);
        return rootP==rootQ;
    }

}
